"""
网络设备开局脚本生成工具
基于Python和Jinja2的多厂商设备配置脚本生成器
Author: kaihonggg
Date: 2025-09-15
"""

import pandas as pd
from jinja2 import Environment, FileSystemLoader
import os
import logging
from datetime import datetime

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('script_generation.log', encoding='utf-8'),
        logging.StreamHandler()
    ]
)

class CommandGenerator:
    def __init__(self, template_dir='templates', output_dir='output'):
        self.template_dir = template_dir
        self.output_dir = output_dir
        # 确保模板目录和输出目录存在
        os.makedirs(template_dir, exist_ok=True)
        os.makedirs(output_dir, exist_ok=True)
        # 设置Jinja2环境
        self.env = Environment(loader=FileSystemLoader(template_dir))
        # 添加全局变量和函数
        self.env.globals['now'] = datetime.now
        
        # 支持的厂商和对应模板
        self.supported_manufacturers = {
            'h3c': 'h3c_template.j2',
            'huawei': 'huawei_template.j2',
            'cisco': 'cisco_template.j2',
            'ruijie': 'ruijie_template.j2'
        }
    
    def read_data(self, data_file):
        """读取数据文件（支持Excel和CSV）"""
        if not os.path.exists(data_file):
            logging.error(f'找不到文件：{data_file}')
            raise FileNotFoundError(f'找不到文件：{data_file}')
            
        logging.info(f'正在读取文件：{data_file}')
        try:
            # 根据文件扩展名选择读取方法
            if data_file.endswith(('.xlsx', '.xls')):
                data = pd.read_excel(data_file)
            else:
                # 默认以CSV格式读取
                data = pd.read_csv(data_file)
                
            if data.empty:
                raise pd.errors.EmptyDataError('文件中没有数据')
                
            logging.info(f'成功读取数据，共 {len(data)} 行')
            return data
        except Exception as e:
            logging.error(f'读取文件时出错：{str(e)}')
            raise
    
    def _get_column_mapping(self, df):
        """获取列名映射，支持标准字段名和中文别名"""
        standard_mapping = {
            'manufacturer': ['manufacturer', 'vendor', '厂商', '品牌', '制造商'],
            'sysname': ['sysname', '设备名称', '设备名', 'hostname', '主机名'],
            'mgmt_vlan': ['mgmt_vlan', '管理vlan', '管理VLAN', 'vlan_id', 'vlan', 'VLAN'],
            'mgmt_ip': ['mgmt_ip', '管理IP', '管理地址', 'ip_address', 'ip', 'IP地址'],
            'mgmt_mask': ['mgmt_mask', '管理掩码', '管理子网掩码', 'subnet_mask', 'mask', '子网掩码'],
            'gateway': ['gateway', '网关', '默认网关'],
            'user': ['user', '用户名', '账户', '账号'],
            'password': ['password', '密码', 'passwd', 'pwd']
        }
        
        column_mapping = {}
        df_columns_lower = {col.lower(): col for col in df.columns}
        
        for standard_name, possible_names in standard_mapping.items():
            for possible_name in possible_names:
                if possible_name.lower() in df_columns_lower:
                    column_mapping[standard_name] = df_columns_lower[possible_name.lower()]
                    break
        
        return column_mapping

    def generate_commands(self, template_name, data_file):
        """生成命令，支持多厂商模板"""
        data = self.read_data(data_file)
        
        # 获取列名映射
        column_mapping = self._get_column_mapping(data)
        
        # 检查必需字段
        required_fields = ['sysname']
        missing_fields = [f for f in required_fields if f not in column_mapping]
        if missing_fields:
            raise ValueError(f"缺少必需字段：{missing_fields}")
        
        # 重命名列以便统一处理
        data = data.rename(columns={v: k for k, v in column_mapping.items()})
        
        # 按sysname分组生成命令
        grouped_commands = {}
        
        # 处理manufacturer列可能为空的情况
        if 'manufacturer' in data.columns:
            # 将空值替换为'default'，确保所有设备都能被处理
            data['manufacturer'] = data['manufacturer'].fillna('default')
            manufacturer_groups = data.groupby('manufacturer')
        else:
            manufacturer_groups = [('default', data)]
        
        for current_manufacturer, group_data in manufacturer_groups:
            # 确定模板
            current_template = template_name
            
            # 只有当厂商名称是字符串且在支持列表中时才替换模板
            if isinstance(current_manufacturer, str) and current_manufacturer.lower() != 'default' and current_manufacturer.lower() in self.supported_manufacturers:
                current_template = self.supported_manufacturers[current_manufacturer.lower()]
            
            try:
                template = self.env.get_template(current_template)
            except Exception as e:
                logging.warning(f"无法加载模板 {current_template}，使用默认模板: {str(e)}")
                # 尝试使用command_template.j2作为后备
                try:
                    template = self.env.get_template('command_template.j2')
                except:
                    logging.error("无法加载任何模板，跳过处理")
                    continue
            
            # 为每个设备生成命令
            for _, row in group_data.iterrows():
                params = row.to_dict()
                try:
                    command = template.render(**params)
                    sysname = params['sysname']
                    if pd.notna(sysname) and sysname.strip():
                        if sysname not in grouped_commands:
                            grouped_commands[sysname] = []
                        grouped_commands[sysname].append(command)
                        logging.info(f'已为 {sysname} 生成命令')
                    else:
                        logging.warning(f'跳过空的sysname')
                except Exception as e:
                    logging.error(f'生成命令时出错：{str(e)}')
        
        return grouped_commands

    def save_commands(self, commands, task_name=None):
        """保存命令到文件"""
        output_path = self.output_dir
        if task_name:
            # 清理任务名称中的非法字符
            safe_task_name = "".join(c for c in str(task_name) if c.isalnum() or c in (' ', '-', '_')).rstrip()
            output_path = os.path.join(self.output_dir, safe_task_name)
            os.makedirs(output_path, exist_ok=True)
            logging.info(f'创建任务目录：{output_path}')
        
        saved_files = []

        
        for device_name, command_list in commands.items():
            filename = f'commands_{device_name}.txt'
            filepath = os.path.join(output_path, filename)
            
            try:
                with open(filepath, 'w', encoding='utf-8') as f:
                    for command in command_list:
                        f.write(command)
                        f.write('\n')
                
                saved_files.append(filepath)
                logging.info(f'已保存文件：{filepath}')
            except Exception as e:
                logging.error(f'保存文件 {filepath} 时出错：{str(e)}')
                raise
        
        return saved_files

def main():
    """主函数 - 交互式界面"""
    generator = CommandGenerator()
    
    print("🚀 网络设备开局脚本生成工具")
    print("=" * 50)
    print("💡 提示：如果manufacturer列为空，将使用默认模板处理")
    
    try:
        # 输入任务名称
        task_name = input("请输入任务名称（留空使用默认）：").strip()
        task_name = task_name if task_name else None
        
        # 选择参数文件
        data_file = input("请输入参数文件路径（默认：parameters.csv）：").strip()
        data_file = data_file if data_file else 'parameters.csv'
        
        # 检查文件是否存在
        if not os.path.exists(data_file):
            print(f"❌ 文件不存在：{data_file}")
            return
        
        # 预览数据
        try:
            preview_data = generator.read_data(data_file)
            print(f"\n📊 数据预览（共{len(preview_data)}行）：")
            print(preview_data.head())
            print(f"列名：{list(preview_data.columns)}")
            
            # 检查厂商信息显示
            has_manufacturer_info = False
            for col in ['manufacturer', 'vendor', '厂商']:
                if col in preview_data.columns and not preview_data[col].isna().all():
                    print(f"厂商分布：{dict(preview_data[col].value_counts())}")
                    has_manufacturer_info = True
                    break
                    
            # 自动模板选择逻辑
            if has_manufacturer_info:
                # 获取有效的厂商数据
                valid_manufacturers = None
                for col in ['manufacturer', 'vendor', '厂商']:
                    if col in preview_data.columns:
                        valid_manufacturers = preview_data[col].dropna().str.lower()
                        if len(valid_manufacturers) > 0:
                            break
                
                if valid_manufacturers is not None and valid_manufacturers.isin(generator.supported_manufacturers.keys()).any():
                    print("\n🤖 检测到厂商列，将根据厂商自动选择模板")
                    print(f"📋 厂商列表：{list(valid_manufacturers.unique())}")
                    # 使用默认模板作为后备，实际会在generate_commands中根据厂商自动选择
                    template_name = 'command_template.j2'
                else:
                    # 没有有效的厂商数据，让用户选择模板
                    template_name = _select_template(generator)
            else:
                # 没有厂商列，让用户选择模板
                template_name = _select_template(generator)
        except Exception as e:
            print(f"❌ 读取文件失败：{e}")
            return
        
        # 生成命令
        print("\n⚙️ 正在生成命令...")
        commands = generator.generate_commands(template_name, data_file)
        
        # 保存命令
        saved_files = generator.save_commands(commands, task_name)
        
        # 显示结果
        print(f"\n✅ 命令生成完成！")
        print(f"📂 共生成 {len(commands)} 个设备的配置文件")
        if commands:
            print(f"💾 文件保存在：{os.path.dirname(saved_files[0]) if saved_files else generator.output_dir}")
            
            # 显示生成的文件列表
            print("\n📋 生成的文件：")
            for file in saved_files:
                print(f"   - {os.path.basename(file)}")
        else:
            print("⚠️  没有生成任何配置文件，请检查参数文件中的数据")
            print("   - 确保sysname列有值")
            print("   - 检查数据格式是否正确")
        
        # 是否显示命令内容
        show_content = input("\n是否显示生成的命令内容？(y/n)：").lower()
        if show_content == 'y':
            for device_name, cmd_list in commands.items():
                print(f"\n📱 {device_name} 配置：")
                print("-" * 30)
                for cmd in cmd_list:
                    print(cmd)
                print()
    
    except KeyboardInterrupt:
        print("\n\n👋 用户中断操作")
    except Exception as e:
        print(f"❌ 发生错误：{e}")
        logging.error(f"主程序错误：{e}")

def _select_template(generator):
    """选择模板辅助函数"""
    template_files = [f for f in os.listdir(generator.template_dir) if f.endswith('.j2')]
    
    # 确保command_template.j2是第1个选项
    if 'command_template.j2' in template_files:
        template_files.remove('command_template.j2')
        template_files.insert(0, 'command_template.j2')
    
    template_mapping = {str(i+1): template for i, template in enumerate(template_files)}
    
    print("\n📁 可用模板：")
    for num, template in template_mapping.items():
        print(f"  {num}. {template}")
    
    choice = input("请选择模板编号（默认：1）：").strip()
    if not choice or choice not in template_mapping:
        return template_mapping.get('1', 'command_template.j2')
    else:
        return template_mapping[choice]

if __name__ == '__main__':
    main()